"""
filename: _exploit_generation.py
author: ww9210
This plugin is used to reproduce the gadget chain and generate concrete payload discovered in previous phase.
We do not generate concrete payload during exploration because such constraint solving will eat a lot of memeoy space
This file does not handle the detail of payload generation, instead, those implementations are in _payload_generation.py
"""
from IPython import embed
import angr
from multiprocessing import Process
from os import listdir
from os.path import isfile, join


class ExploitGenererationMixin:
    def parse_gadget_chain_log(self, gadget_chain_log):
        """
        get gadget from the log file
        :param gadget_chain_log:
        :return:
        """
        '''
        self.disclosure_gadgets = pickle.load(open(disclosure_gadget_path, 'rb'))
        self.fake_stack_gadgets = pickle.load(open(fake_stack_gadget_path, 'rb'))
        self.smash_gadgets = pickle.load(open(smash_gadget_path, 'rb'))
        self.bloom_gadgets = pickle.load(open(bloom_gadget_path, 'rb'))
        self.fork_gadgets = pickle.load(open(fork_gadget_path, 'rb'))
        '''
        with open(gadget_chain_log,'r') as f:
            log = f.readlines()
        '''
        blooming gadget:event_sched_out
        forking gadget:em_call
        prologue gadget:tg3_get_5720_nvram_info
        disclosure gadget:mmc_ioctl_dvd_auth
        smash gadget:fb_set_user_cmap
        '''
        bloom_gadget_name = log[0].split(':')[1].strip('\n')
        forking_gadget_name = log[1].split(':')[1].strip('\n')
        prologue_gadget_name = log[2].split(':')[1].strip('\n')
        disclosure_gadget_name = log[3].split(':')[1].strip('\n')
        smash_gadget_name = log[4].split(':')[1].strip('\n')
        timestamp = int(log[5].strip('\n'), 16)
        self.current_timestamp = timestamp
        good_bloom_gadget_idx = -1
        good_forking_gadget_idx = -1
        good_prologue_disclosure_pair_idx = -1
        good_smash_gadget_idx = -1
        for i, bloom_gadget in enumerate(self.bloom_gadgets):
            if bloom_gadget[1] == bloom_gadget_name:
                good_bloom_gadget_idx = i
                self.current_bloom_gadget_to_reproduce = bloom_gadget  # set current bloom gadget to reproduce
                break
        if good_bloom_gadget_idx == -1:
            assert 0  # wtf

        # get forking
        for i, forking_gadget in enumerate(self.fork_gadgets):
            if forking_gadget[1] == forking_gadget_name:
                good_forking_gadget_idx = i
                self.current_forking_gadget_to_reproduce = forking_gadget  # set current forking gadget to reproduce
                break
        if good_forking_gadget_idx == -1:
            print(forking_gadget_name, forking_gadget_name.encode('hex'))
            assert 0  # wtf

        # get good prologue disclosure pair to reproduce from the log
        self.current_prologue_disclosure_pair_to_reproduce = []
        for i, prologue_disclosure_pair in enumerate(self.prologue_disclosure_pairs):
            prologue_gadget = prologue_disclosure_pair[0]
            disclosure_gadget = prologue_disclosure_pair[1]
            if prologue_gadget[5] == prologue_gadget_name and disclosure_gadget[3] == disclosure_gadget_name:
                good_prologue_disclosure_pair_idx = i
                self.current_prologue_disclosure_pair_to_reproduce.append(prologue_disclosure_pair)
        if good_prologue_disclosure_pair_idx == -1:
            assert 0

        # get current smash gadget to reproduce
        self.current_smash_gadget_to_reproduce = []
        for i, smash_gadget in enumerate(self.smash_gadgets):
            if smash_gadget[3] == smash_gadget_name:
                good_smash_gadget_idx = i
                self.current_smash_gadget_to_reproduce.append(smash_gadget)  # set current smash gadget to reproduce
        if good_smash_gadget_idx == -1:
            assert 0  # wtf

        return

    def update_gadget_candidate_list(self):
        """
        replace gadget candidate list
        :return:
        """
        self.bloom_gadgets = []
        self.fork_gadgets = []
        self.prologue_disclosure_pairs = []
        self.smash_gadgets = []

        # put gadget chain to reproduce to candicate list
        self.bloom_gadgets.append(self.current_bloom_gadget_to_reproduce)
        self.fork_gadgets.append(self.current_forking_gadget_to_reproduce)
        for pair in self.current_prologue_disclosure_pair_to_reproduce:
            self.prologue_disclosure_pairs.append(pair)
        for smash_gadget in self.current_smash_gadget_to_reproduce:
            self.smash_gadgets.append(smash_gadget)

        return

    def update_generated_payload_number(self):
        """
        update num of generated payload so we won't overwrite generated payloads
        :return:
        """
        payload_path = self.payload_path
        mypath = payload_path
        onlyfiles = [f for f in listdir(mypath) if isfile(join(mypath, f))]
        f_cnt = 0
        for fn in onlyfiles:
            if 'info_' in fn:
                f_cnt += 1
        self.num_of_generate_payload = f_cnt

    def reproduce_an_exploit(self, gadget_chain_log_filename=None, debug_dfs=False):
        if gadget_chain_log_filename is None:
            return

        # get gadget chain to reproduce from log
        self.parse_gadget_chain_log(gadget_chain_log_filename)

        # replace gadget list to test
        self.update_gadget_candidate_list()
        # set several global variables
        self.current_log_to_reproduce = gadget_chain_log_filename
        self.is_dfs_search_routine = True
        self.debug_dfs = debug_dfs
        self.reproduce_mode = True

        # start reproducing
        self._bloom(start_bloom_idx=0)
        return

    def reproduce_all_log_in_a_dir(self, number_to_reproduce=None, use_subprocess=True, gadget_chain_log_dir=None):
        """
        :param gadget_chain_log_dir: path of generated logs to reproduce
        :return:
        """
        # firstly update the total number of payloads we already reproduced

        # respect the flow of osok_dfs, but not necessary here
        self.prologue_disclosure_pairs = self.get_prologue_disclosure_pairs()

        # load qemu snapshot
        self.load_hyper_parameters()
        self.load_qemu_snapshot()
        self.initial_state = self.get_initial_state(control_memory_base=self.controlled_memory_base)

        if gadget_chain_log_dir is None:
            return
        mypath = gadget_chain_log_dir
        onlyfiles = [f for f in listdir(mypath) if isfile(join(mypath, f))]
        files_with_full_path = [join(mypath, f) for f in listdir(mypath)]
        for fn in files_with_full_path:
            self.update_generated_payload_number()
            if 'info_' in fn:
                print('[+] generating exploit for', fn)
                if use_subprocess:
                    p = Process(target=self.reproduce_an_exploit, args=(fn, False))
                    p.start()
                    p.join()
                else:
                    self.reproduce_an_exploit(fn, False)
        return
